import os
import utils_graph
import torch
import torch.optim as optim
from models import H_GAT
import torch.nn as nn
from freq_statistic import fre_statis
import torch.utils.data as Data

os.environ['CUDA_VISIBLE_DEVICES'] = '-1'  # 不用gpu，cuda有点问题

train_feat, train_label, train_wei, test_feat, test_label, test_wei = utils_graph.load_data()
batch_size = 16
dataset = Data.TensorDataset(train_wei, train_feat, train_label)
train_loader = Data.DataLoader(dataset=dataset, batch_size=batch_size, shuffle=True)

model = H_GAT(70, 80, 2)

n = 5
LR = 0.00001
path_threshold = 0.5
EPOCH = 101
max_acc = 0
loss_list = []
acc_list = []
out_data = torch.zeros(420, 2)

optimizer = optim.Adam(model.parameters(), lr=LR)  # 优化器
loss_func = nn.CrossEntropyLoss()
for epoch in range(EPOCH):
    for step, (b_w, b_feat, b_lab) in enumerate(train_loader):
        output = model.forward(b_w,b_feat,n,path_threshold)
        loss = loss_func(output, b_lab)
        acc_val = utils_graph.accuracy(output, b_lab)
        optimizer.zero_grad()  # 清空梯度
        loss.backward()  # 反向传播
        optimizer.step()

        if epoch == EPOCH - 1:
            if output.shape[0] == batch_size:
                out_data[step * 16:(step + 1) * output.shape[0], :] = output

    if epoch % 10 == 0:
        print('Epoch: {:04d}'.format(epoch + 1), 'loss_train: {:.4f}'.format(loss.item()),
              'acc_val: {:.4f}'.format(acc_val.item()))

        model.eval()
        output1 = model(test_wei, test_feat, n, path_threshold)
        loss_val1 = nn.CrossEntropyLoss()(output1, test_label)
        acc_val1 = utils_graph.accuracy(output1, test_label)

        print("~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~")
        print("Test set results:",
              "loss= {:.4f}".format(loss_val1.item()),
              "accuracy= {:.4f}".format(acc_val1.item()))
        print("~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~")
        loss_list.append(float(loss_val1.item()))
        acc_list.append(float(acc_val1.item()))
    if max_acc < acc_val1:
        max_acc = acc_val1
        TP, TN, FN, FP = utils_graph.stastic_indicators(output1, test_label)
        ACC = (TP + TN) / (TP + TN + FP + FN)
        SEN = TP / (TP + FN)
        SPE = TN / (FP + TN)
        BAC = (SEN + SPE) / 2
        output2 = output1

# 特征提取
fc1_w = model.state_dict()['fc1.weight']
fc2_w = model.state_dict()['fc2.weight']
fc3_w = model.state_dict()['fc3.weight']
fc4 = out_data
best_imp, best_idx = fre_statis(fc1_w, fc2_w, fc3_w, fc4)
for num in range(best_idx.shape[0]):
    if fc4[num, 0] < fc4[num, 1]:
        best_index = best_idx[num]
        best_important = best_imp[num]
        break